from pyro.nn import PyroSample, PyroModule
from pyro import distributions as dist

class Location(PyroModule):
    def __init__(self):
        super().__init__()
        # Independent priors
        self.radius = PyroSample(dist.LogNormal(0,1))
        self.theta = PyroSample(dist.Normal(0.5, 0.1))
        # Dependent deterministic
        self.true_x = PyroSample(lambda self: self.radius * self.theta.cos())
        self.true_y = PyroSample(lambda self: self.radius * self.theta.sin())
        # Dependent samples
        self.observed_x = PyroSample(lambda self: dist.Normal(self.true_x, 0.05))
        self.observed_y = PyroSample(lambda self: dist.Normal(self.true_y, 0.05))
        
    def forward(self):
        return self.true_x, self.true_y, self.observed_x, self.observed_y